#include "kwdb/client.h"
#include "kwdb/stream_client.h"
#include "kwdb/batcher.h"
#include "kwdb/retry_policy.h"
#include "kwdb/metrics.h"
#include <iostream>
#include <chrono>
#include <thread>
#include <atomic>
#include <vector>
#include <cstring>
#include <random>
#include <iomanip>
#include <cstdlib>

std::atomic<int64_t> g_success_count{0};
std::atomic<int64_t> g_error_count{0};

// Generate random string
std::string random_string(int len) {
    static const char charset[] = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789";
    static thread_local std::mt19937 gen(std::random_device{}());
    static std::uniform_int_distribution<> dist(0, sizeof(charset) - 2);

    std::string result;
    result.reserve(len);
    for (int i = 0; i < len; ++i) {
        result += charset[dist(gen)];
    }
    return result;
}

void worker_thread(int thread_id, int total_per_thread, int batch_size) {
    kwdb::IngestConfig config;

    // Read from environment or use defaults
    const char* endpoint = std::getenv("INGEST_ENDPOINT");
    if (endpoint) {
        config.ingest_endpoint = endpoint;
    }

    kwdb::KwdbClient client(config);

    std::mt19937 gen(thread_id);
    std::uniform_real_distribution<> value_dist(0.0, 100.0);

    // Use monotonic timestamp: start from current time + thread_id offset, increment by 1us per point
    auto base_time = std::chrono::system_clock::now() + std::chrono::seconds(thread_id * 10000);  // Current time + offset per thread
    int64_t point_counter = 0;

    int written = 0;
    while (written < total_per_thread) {
        std::vector<kwdb::Point> batch;
        int batch_count = std::min(batch_size, total_per_thread - written);

        for (int i = 0; i < batch_count; ++i) {
            kwdb::Point p("sdk_bench");
            p.add_tag("host", "bench-host-" + std::to_string(thread_id % 4));
            p.add_tag("region", "bench-region-" + std::to_string(thread_id % 2));
            p.add_field("usage", value_dist(gen));
            p.add_field("temperature", value_dist(gen));

            // Monotonic timestamp: 1us increment per point
            auto ts = base_time + std::chrono::microseconds(point_counter);
            p.timestamp(ts);
            point_counter++;

            batch.push_back(std::move(p));
        }

        kwdb::Status status = client.batch_write(batch);
        if (status.ok()) {
            g_success_count += batch_count;
        } else {
            g_error_count += batch_count;
            std::cerr << "[thread-" << thread_id << "] Error: "
                      << status.message << std::endl;
        }

        written += batch_count;
    }
}

int main(int argc, char** argv) {
    if (argc < 4) {
        std::cout << "Usage: " << argv[0] << " <total_points> <batch_size> <concurrency>" << std::endl;
        std::cout << "Example: " << argv[0] << " 100000 200 8" << std::endl;
        std::cout << std::endl;
        std::cout << "Environment variables:" << std::endl;
        std::cout << "  INGEST_ENDPOINT - ingestd address (default: 127.0.0.1:9090)" << std::endl;
        return 1;
    }

    int total_points = std::atoi(argv[1]);
    int batch_size = std::atoi(argv[2]);
    int concurrency = std::atoi(argv[3]);

    if (total_points <= 0 || batch_size <= 0 || concurrency <= 0) {
        std::cerr << "Error: All parameters must be positive integers" << std::endl;
        return 1;
    }

    int points_per_thread = total_points / concurrency;
    int actual_total = points_per_thread * concurrency;

    std::cout << "=== SDK Benchmark ===" << std::endl;
    std::cout << "Total points:  " << actual_total << std::endl;
    std::cout << "Batch size:    " << batch_size << std::endl;
    std::cout << "Concurrency:   " << concurrency << std::endl;
    std::cout << "Per thread:    " << points_per_thread << std::endl;
    std::cout << "Endpoint:      " << (std::getenv("INGEST_ENDPOINT")
                                        ? std::getenv("INGEST_ENDPOINT")
                                        : "127.0.0.1:9090") << std::endl;
    std::cout << std::endl;

    auto start = std::chrono::steady_clock::now();

    std::vector<std::thread> threads;
    for (int i = 0; i < concurrency; ++i) {
        threads.emplace_back(worker_thread, i, points_per_thread, batch_size);
    }

    for (auto& t : threads) {
        t.join();
    }

    auto end = std::chrono::steady_clock::now();
    auto duration_ms = std::chrono::duration_cast<std::chrono::milliseconds>(end - start).count();
    double duration_sec = duration_ms / 1000.0;

    int64_t success = g_success_count.load();
    int64_t errors = g_error_count.load();
    double qps = success / duration_sec;

    std::cout << "=== Results ===" << std::endl;
    std::cout << "Duration:      " << std::fixed << std::setprecision(2)
              << duration_sec << " seconds" << std::endl;
    std::cout << "Success:       " << success << " points" << std::endl;
    std::cout << "Errors:        " << errors << " points" << std::endl;
    std::cout << "QPS:           " << std::fixed << std::setprecision(2)
              << qps << " points/sec" << std::endl;
    std::cout << "Avg latency:   " << std::fixed << std::setprecision(2)
              << (duration_ms * concurrency / (double)(success/batch_size))
              << " ms/batch" << std::endl;
    std::cout << std::endl;

    if (errors > 0) {
        std::cout << "� Warning: " << errors << " points failed to write" << std::endl;
        return 1;
    }

    std::cout << " Benchmark completed successfully" << std::endl;
    return 0;
}
